from skl2onnx.common.data_types import FloatTensorType
from onnxmltools import convert_lightgbm,convert_sklearn,convert_xgboost
# from skl2onnx import convert_sklearn
from lightgbm import LGBMRegressor
from xgboost import XGBRegressor
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.tree import ExtraTreeRegressor,DecisionTreeRegressor
import numpy as np

class MLPriceModel:
    def __init__(self, model:str=['XGB','LGB','GDBT','MLP','ET','DT']) -> None:
        if model == 'XGB':
            self.model = XGBRegressor(device='gpu')
            self.NAME = 'XGBRegressor'
        elif model == 'LGB':
            self.model = LGBMRegressor(device='gpu')
            self.NAME = 'LGBMRegressor'
        elif model == 'GDBT':
            self.model = GradientBoostingRegressor(loss='absolute_error')
            self.NAME = 'GradientBoostingRegressor'
        elif model == 'MLP':
            self.model = MLPRegressor()
            self.NAME = 'MLPRegressor'
        elif model == 'ET':
            self.model = ExtraTreeRegressor(criterion='absolute_error')
            self.NAME = 'ExtraTreeRegressor'
        elif model == 'DT':
            self.model = DecisionTreeRegressor(criterion='absolute_error')
            self.NAME = 'DecisionTreeRegressor'
        
        # self.model = MLPRegressor()
    def fit(self, data:np.ndarray, label:np.ndarray, **args):
        '''
        用于训练模型
        '''
        self.model.fit(data,label)
        
    def predict(self, data:np.ndarray, **args) -> np.ndarray:
        '''
        对数据进行预测
        '''
        out = self.model.predict(data)
        
        return np.expand_dims(out,-1)
    
    def save_to_onnx(self, feature_size:int,path:str):
        initial_type = [('float_input', FloatTensorType([None, feature_size]))]
        if isinstance(self.model, XGBRegressor):
            onnx_model = convert_xgboost(self.model,initial_types=initial_type)
        elif isinstance(self.model, LGBMRegressor):
            onnx_model = convert_lightgbm(self.model,initial_types=initial_type)
        else:
            onnx_model = convert_sklearn(self.model,initial_types=initial_type)
        with open(path,'bw') as f:
            f.write(onnx_model.SerializeToString())
    
    def get_name(self):
        return self.NAME